import pyaudio
import wave
import time
import os
import threading
from datetime import datetime
import logging
import keyboard
import numpy as np
import librosa

# 配置参数
FORMAT = pyaudio.paInt16
CHANNELS = 1
RATE = 44100
CHUNK = 1024
OUTPUT_DIR = "audio_recordings"
PCM_OUTPUT_DIR = "pcm_data"  # 新增：PCM 数据保存目录


# 创建目录
os.makedirs(OUTPUT_DIR, exist_ok=True)
os.makedirs(PCM_OUTPUT_DIR, exist_ok=True)  # 创建 PCM 数据保存目录

# 初始化 PyAudio
p = pyaudio.PyAudio()

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

def wav_to_pcm(wav_filepath, sample_rate=16000, channels=1):
    """
    将 WAV 文件转换为 PCM 数据 (bytes)。

    Args:
        wav_file_path: WAV 文件的路径。

    Returns:
        PCM 数据 (bytes)，如果转换失败则返回 None。
    """
    """将WAV文件转换为PCM格式"""
    try:
        y, sr = librosa.load(wav_filepath, sr=sample_rate, mono=(channels == 1))

        logging.debug(f"Loaded audio data: shape={y.shape}, dtype={y.dtype}")
        logging.debug(f"Loaded audio data: max={np.max(y)}, min={np.min(y)}")

        if len(y.shape) > 1 and channels == 1:
            y = librosa.to_mono(y)
            logging.debug("Converted to mono")
        elif len(y.shape) == 1 and channels == 2:
            y = np.stack([y, y], axis=0)
            logging.debug("Stacked to stereo")

        y = (y * 32767).astype(np.int16)
        logging.debug(f"Converted to int16: max={np.max(y)}, min={np.min(y)}")
        return y.tobytes()
    except Exception as e:
        logging.error(f"WAV转PCM失败: {wav_filepath}, 错误: {e}")
        return None

def record_audio(duration=None):
    """
    录制音频，空格键开始/结束，返回 PCM 数据。

    Args:
        duration: 录音时长 (秒)，如果为 None，则通过空格键控制。

    Returns:
        PCM 数据 (bytes)，如果录音失败则返回 None。
    """

    frames = []  # 存储音频帧
    is_recording = False
    is_endding = True
    stream = None
    pcm_data = None  # 初始化 pcm_data
    stop_event = threading.Event()  # 使用事件来停止循环
    def _record():
        nonlocal frames, is_recording, is_endding, stream, pcm_data

        try:
            stream = p.open(format=FORMAT, channels=CHANNELS, rate=RATE, input=True, frames_per_buffer=CHUNK)

            if duration:
                start_time = time.time()
                while time.time() - start_time < duration:
                    try:
                        data = stream.read(CHUNK, exception_on_overflow=False)  # 确保有数据
                        frames.clear()
                        frames.append(data)
                    except Exception as e:
                        logging.error(f"读取音频流失败: {e}")
                        break  # 如果读取失败，则跳出循环
            else:
                while not stop_event.is_set():  # 使用事件来控制循环
                    try:
                        data = stream.read(CHUNK, exception_on_overflow=False)  # 确保有数据
                        frames.append(data)
                    except Exception as e:
                        logging.error(f"读取音频流失败: {e}")
                        break  # 如果读取失败，则跳出循环

        except Exception as e:
            logging.error(f"录音线程异常: {e}")
        finally:
            if stream:
                stream.stop_stream()
                stream.close()


    try:
        if duration is None:
            # 使用空格键控制
            def start_stop():
                nonlocal is_recording, is_endding, frames, pcm_data

            # 使用空格键控制
                if not is_recording:
                    is_recording = True
                    frames = []  # 重置 frames

                    recording_thread = threading.Thread(target=_record)
                    recording_thread.daemon = True
                    recording_thread.start()
                    logging.info("开始录音")
                else:
                    is_recording = False
                    is_endding = False  # 停止录音时设为True，表示处理结束
                    stop_event.set() # 设置事件，停止录音循环
                    logging.info("停止录音")
                    #  将 frames 转换为 PCM 数据
                    if frames:
                        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                        wav_file_path = os.path.join(OUTPUT_DIR, f"recording_{timestamp}.wav")
                        try:
                            # 先保存为WAV
                            with wave.open(wav_file_path, 'wb') as wf:
                                wf.setnchannels(CHANNELS)
                                wf.setsampwidth(p.get_sample_size(FORMAT))
                                wf.setframerate(RATE)
                                wf.writeframes(b''.join(frames))
                            #logging.info(f"WAV 文件已保存到 {wav_file_path}")

                            # 将 WAV 转换为 PCM
                            pcm_data = wav_to_pcm(wav_file_path)  # 调用新函数
                            if pcm_data is None:
                                logging.error("从 WAV 文件转换到 PCM 失败")

                        except Exception as e:
                            logging.error(f"保存 WAV 文件或读取PCM数据失败: {e}")
                            pcm_data = None  # 如果失败，返回None


                    else:

                        logging.warning("没有录制到任何音频数据。")


            keyboard.add_hotkey('space', start_stop)
            print("按空格键开始/停止录音")
            # 阻塞等待录音结束
            while is_recording or is_endding:
                time.sleep(0.1)

        else:
            # 固定时长录音
            logging.info("开始录音 (固定时长)")
            is_recording = True
            _record()
            is_recording = False
            logging.info("停止录音 (固定时长)")


    except KeyboardInterrupt:
        logging.info("程序被用户中断")
    except Exception as e:
        logging.error(f"发生错误: {e}")
        return None  # 录音失败返回 None
    return pcm_data  # 返回 PCM 数据
if __name__ == '__main__':
    # 调用示例
    pcm_data = record_audio()  # 空格键控制
    if pcm_data:
        try:
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            wav_file_path = os.path.join(OUTPUT_DIR, f"playback_{timestamp}.wav")

            with wave.open(wav_file_path, 'wb') as wf:
                wf.setnchannels(CHANNELS)
                wf.setsampwidth(p.get_sample_size(FORMAT))
                wf.setframerate(RATE)
                wf.writeframes(pcm_data)  # 直接写入 PCM 数据
            print(f"PCM 数据已转换为 WAV 并保存到: {wav_file_path}")
        except Exception as e:
            print(f"创建 WAV 文件时出错: {e}")
    else:
        print("录音失败")